#!/usr/bin/env python3
"""同步方式检查 ChromaDB 数据"""

import chromadb
import os

def check_chromadb_sync():
    """同步方式检查 ChromaDB 数据"""
    print("=== 同步方式检查 ChromaDB 数据 ===")
    
    try:
        # 连接到 ChromaDB
        chroma_path = "./data/chroma_db"
        print(f"连接到 ChromaDB: {chroma_path}")
        
        client = chromadb.PersistentClient(path=chroma_path)
        
        # 列出所有集合
        collections = client.list_collections()
        print(f"\n找到 {len(collections)} 个集合:")
        for collection in collections:
            print(f"  - {collection.name}")
        
        # 获取 conversations 集合
        try:
            conversations_collection = client.get_collection("conversations")
            print(f"\n✅ 成功获取 conversations 集合")
        except Exception as e:
            print(f"❌ 获取 conversations 集合失败: {e}")
            return
        
        # 获取集合中的所有数据
        print(f"\n1. 获取集合中的所有数据...")
        all_data = conversations_collection.get()
        
        print(f"   总文档数: {len(all_data['ids'])}")
        
        if len(all_data['ids']) == 0:
            print(f"   ❌ 集合为空")
            return
        
        # 分析前几个文档
        print(f"\n2. 分析前几个文档...")
        
        for i in range(min(len(all_data['ids']), 5)):
            print(f"\n   文档 [{i+1}]:")
            print(f"     ID: {all_data['ids'][i]}")
            
            # 内容
            content = all_data['documents'][i] if all_data['documents'] else "N/A"
            print(f"     内容: {content[:100]}...")
            
            # 元数据
            metadata = all_data['metadatas'][i] if all_data['metadatas'] else {}
            print(f"     元数据:")
            for key, value in metadata.items():
                print(f"       {key}: {value}")
        
        # 统计用户ID和对话ID
        print(f"\n3. 统计用户ID和对话ID...")
        user_ids = {}
        conversation_ids = set()
        
        for i in range(len(all_data['ids'])):
            metadata = all_data['metadatas'][i] if all_data['metadatas'] else {}
            user_id = metadata.get('user_id', 'unknown')
            conv_id = metadata.get('conversation_id', 'unknown')
            
            if user_id not in user_ids:
                user_ids[user_id] = 0
            user_ids[user_id] += 1
            
            conversation_ids.add(conv_id)
        
        print(f"   用户ID分布:")
        for user_id, count in user_ids.items():
            print(f"     {user_id}: {count} 条消息")
        
        print(f"\n   对话ID列表:")
        for conv_id in sorted(conversation_ids):
            print(f"     {conv_id}")
        
        # 搜索测试
        print(f"\n4. 搜索测试...")
        
        # 不带用户过滤的搜索
        print(f"   搜索 '王者荣耀' (无过滤):")
        try:
            results_no_filter = conversations_collection.query(
                query_texts=["王者荣耀"],
                n_results=10
            )
            
            print(f"     结果数量: {len(results_no_filter['ids'][0])}")
            for i, result_id in enumerate(results_no_filter['ids'][0]):
                metadata = results_no_filter['metadatas'][0][i]
                distance = results_no_filter['distances'][0][i]
                print(f"     [{i+1}] ID: {result_id}")
                print(f"         对话ID: {metadata.get('conversation_id')}")
                print(f"         用户ID: {metadata.get('user_id')}")
                print(f"         距离: {distance}")
        except Exception as e:
            print(f"     ❌ 搜索失败: {e}")
        
        # 带用户过滤的搜索
        print(f"\n   搜索 '王者荣耀' (user_id=default_user):")
        try:
            results_with_filter = conversations_collection.query(
                query_texts=["王者荣耀"],
                n_results=10,
                where={"user_id": "default_user"}
            )
            
            print(f"     结果数量: {len(results_with_filter['ids'][0])}")
            for i, result_id in enumerate(results_with_filter['ids'][0]):
                metadata = results_with_filter['metadatas'][0][i]
                distance = results_with_filter['distances'][0][i]
                print(f"     [{i+1}] ID: {result_id}")
                print(f"         对话ID: {metadata.get('conversation_id')}")
                print(f"         用户ID: {metadata.get('user_id')}")
                print(f"         距离: {distance}")
        except Exception as e:
            print(f"     ❌ 带过滤搜索失败: {e}")
        
        print(f"\n✅ 检查完成")
        
    except Exception as e:
        print(f"❌ 检查失败: {e}")
        import traceback
        traceback.print_exc()

if __name__ == "__main__":
    check_chromadb_sync()